import torch
import torch.nn as nn
import math
from torch.nn.utils.rnn import pad_sequence

# ----------------------------
# 1. 1D CNN 编码器（降采样）
# ----------------------------
class CNNEncoder(nn.Module):
    """
    1D CNN 编码器，用于对输入序列进行降采样和特征提取。
    
    参数:
        input_dim (int): 输入特征维度，默认为 8（传感器数据的特征数）。
        hidden_dim (int): 中间层特征维度，默认为 64。
        output_dim (int): 输出特征维度，默认为 128。
    
    输入张量形状:
        x: (B, T, input_dim)，其中 B 是批次大小，T 是时间步长。
    
    输出张量形状:
        (B, T/4, output_dim)，经过两次降采样后时间步长变为 T/4。
    """
    def __init__(self, input_dim=8, hidden_dim=64, output_dim=128):
        super().__init__()
        # 第一层卷积：输入维度 → hidden_dim，时间步长 T → T/2
        self.conv1 = nn.Conv1d(input_dim, hidden_dim, kernel_size=7, stride=2, padding=3)
        # 第二层卷积：hidden_dim → output_dim，时间步长 T/2 → T/4
        self.conv2 = nn.Conv1d(hidden_dim, output_dim, kernel_size=5, stride=2, padding=2)
        self.relu = nn.ReLU()
        self.output_dim = output_dim

    def forward(self, x):
        # 转置输入张量：(B, T, input_dim) → (B, input_dim, T)
        x = x.transpose(1, 2)
        # 第一层卷积 + ReLU：(B, input_dim, T) → (B, hidden_dim, T/2)
        x = self.relu(self.conv1(x))
        # 第二层卷积 + ReLU：(B, hidden_dim, T/2) → (B, output_dim, T/4)
        x = self.relu(self.conv2(x))
        # 转置输出张量：(B, output_dim, T/4) → (B, T/4, output_dim)
        return x.transpose(1, 2)

# ----------------------------
# 2. 位置编码
# ----------------------------
class PositionalEncoding(nn.Module):
    """
    位置编码模块，用于为序列数据添加位置信息。
    
    参数:
        d_model (int): 特征维度，必须与输入张量的最后一维一致。
        max_len (int): 支持的最大序列长度，默认为 500。
    
    输入张量形状:
        x: (B, T, d_model)，其中 B 是批次大小，T 是时间步长。
    
    输出张量形状:
        (B, T, d_model)，与输入形状相同，但添加了位置编码信息。
    """
    def __init__(self, d_model, max_len=500):
        super().__init__()
        # 初始化位置编码矩阵：(max_len, d_model)
        pe = torch.zeros(max_len, d_model)
        # 生成位置索引：(max_len, 1)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        # 计算频率项：用于生成正弦和余弦编码
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        # 填充正弦和余弦编码
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数索引：正弦编码
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数索引：余弦编码
        # 扩展维度：(1, max_len, d_model)
        pe = pe.unsqueeze(0)
        # 注册为缓冲区，避免被优化器更新
        self.register_buffer('pe', pe)

    def forward(self, x):
        # 添加位置编码：x + pe[:, :x.size(1), :]
        x = x + self.pe[:, :x.size(1), :]
        return x

# ----------------------------
# 3. CNN-Transformer Seq2Seq
# ----------------------------
class GaitCNNTransformer(nn.Module):
    """
    CNN-Transformer 序列到序列模型，用于处理变长传感器数据并预测目标序列。
    
    参数:
        input_dim (int): 输入特征维度，默认为 8（传感器数据的特征数）。
        output_dim (int): 输出特征维度，默认为 1（预测目标维度）。
        d_model (int): Transformer 特征维度，默认为 128。
        nhead (int): 多头注意力头数，默认为 8。
        num_layers (int): Transformer 层数，默认为 3。
        dim_feedforward (int): 前馈网络维度，默认为 512。
        dropout (float): Dropout 概率，默认为 0.1。
    
    输入张量形状:
        src: (B, T1, input_dim) —— 变长传感器输入序列。
        tgt: (B, T2, output_dim) —— 变长目标序列（训练时使用）。
        src_key_padding_mask: (B, T1') —— 源序列的填充掩码（CNN 降采样后）。
        tgt_key_padding_mask: (B, T2) —— 目标序列的填充掩码。
    
    输出张量形状:
        (B, T2, output_dim) —— 预测结果。
    """
    def __init__(self, input_dim=8, output_dim=1, d_model=128, nhead=8, num_layers=3, dim_feedforward=512, dropout=0.1):
        super().__init__()
        # CNN 编码器：输入 → d_model
        self.cnn_encoder = CNNEncoder(input_dim, hidden_dim=64, output_dim=d_model)
        # 位置编码：添加位置信息
        self.pos_encoder = PositionalEncoding(d_model, max_len=200)
        
        # Transformer 编码器
        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers)
        
        # Transformer 解码器
        decoder_layer = nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, batch_first=True)
        self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers)
        
        # 输出投影：d_model → output_dim
        self.output_proj = nn.Linear(d_model, output_dim)
        
        # 学习的起始 token：用于解码器的初始输入
        self.start_token = nn.Parameter(torch.randn(1, 1, d_model))

    def forward(self, src, tgt=None, src_key_padding_mask=None, tgt_key_padding_mask=None):
        """
        前向传播逻辑：
        1. 使用 CNN 编码器对输入序列进行降采样和特征提取。
        2. 添加位置编码并通过 Transformer 编码器进一步编码。
        3. 使用 Transformer 解码器生成预测结果。
        
        参数:
            src: (B, T1, input_dim) —— 源序列。
            tgt: (B, T2, output_dim) —— 目标序列（训练时使用）。
            src_key_padding_mask: (B, T1') —— 源序列的填充掩码。
            tgt_key_padding_mask: (B, T2) —— 目标序列的填充掩码。
        
        返回:
            (B, T2, output_dim) —— 预测结果。
        """
        # 1. CNN 编码
        memory = self.cnn_encoder(src)  # (B, T1', d_model)
        T1_prime = memory.size(1)
        
        # 2. 位置编码 + Transformer 编码
        memory = self.pos_encoder(memory)
        if src_key_padding_mask is not None:
            # 调整 mask 长度（因 CNN 降采样）
            src_key_padding_mask = src_key_padding_mask[:, ::4]  # 对每个样本（第0维），每隔4个元素取1个，从索引 0 开始，粗略调整，实际需精确计算 
            src_key_padding_mask = src_key_padding_mask[:, :T1_prime]
        memory = self.transformer_encoder(memory, src_key_padding_mask=src_key_padding_mask)
        print('memory.shape after transformer_encoder:',memory.shape)
        # 3. 解码
        B = src.size(0)
        if tgt is not None:
            # 训练：使用 teacher forcing
            T2 = tgt.size(1)
            tgt_embed = self.start_token.expand(B, T2, -1)  # 简化：实际可用 tgt 投影
            tgt_embed = self.pos_encoder(tgt_embed)
            output = self.transformer_decoder(
                tgt_embed, memory,  # (1) 目标序列的嵌入表示 # (2) 编码器输出（上下文）
                tgt_mask=nn.Transformer.generate_square_subsequent_mask(T2).to(src.device),# (3) 解码器自注意力的因果掩码
                tgt_key_padding_mask=tgt_key_padding_mask,  # (4) 目标序列的 padding 掩码
                memory_key_padding_mask=src_key_padding_mask  # (5) 编码器输出的 padding 掩码
            )
        else:
            # 推理：自回归生成（此处简化为固定长度，实际可用 stop token）
            T2 = 120  # 或动态生成
            tgt_embed = self.start_token.expand(B, T2, -1)
            tgt_embed = self.pos_encoder(tgt_embed)
            output = self.transformer_decoder(
                tgt_embed, memory,
                tgt_mask=nn.Transformer.generate_square_subsequent_mask(T2).to(src.device),
                memory_key_padding_mask=src_key_padding_mask
            )

        return self.output_proj(output)  # (B, T2, 1)

# ----------------------------
# 4. 数据处理辅助函数
# ----------------------------
def collate_fn(batch):
    """
    处理变长序列的批处理函数，用于 DataLoader。
    
    参数:
        batch: 包含多个变长序列的列表，每个元素为 (src, tgt) 元组。
    
    返回:
        src_padded: (B, max_T1, 8) —— 填充后的源序列。
        tgt_padded: (B, max_T2, 1) —— 填充后的目标序列。
        src_mask: (B, max_T1) —— 源序列的填充掩码（True 表示填充部分）。
        tgt_mask: (B, max_T2) —— 目标序列的填充掩码（True 表示填充部分）。
    """
    src_list, tgt_list = zip(*batch)
    src_lengths = [len(x) for x in src_list]
    tgt_lengths = [len(y) for y in tgt_list]
    
    # 填充源序列和目标序列
    src_padded = pad_sequence(src_list, batch_first=True, padding_value=0)
    tgt_padded = pad_sequence(tgt_list, batch_first=True, padding_value=0)
    
    # 生成填充掩码
    src_mask = torch.zeros(len(src_list), src_padded.size(1), dtype=torch.bool)
    tgt_mask = torch.zeros(len(tgt_list), tgt_padded.size(1), dtype=torch.bool)
    for i, (src_len, tgt_len) in enumerate(zip(src_lengths, tgt_lengths)):
        src_mask[i, src_len:] = True
        tgt_mask[i, tgt_len:] = True
        
    return src_padded, tgt_padded, src_mask, tgt_mask

# ----------------------------
# 5. 使用示例
# ----------------------------
if __name__ == "__main__":
    """
    示例代码：训练 CNN-Transformer 模型。
    
    步骤:
        1. 生成模拟变长数据。
        2. 创建 DataLoader 并使用 collate_fn 处理批数据。
        3. 初始化模型、损失函数和优化器。
        4. 训练循环：计算损失并更新模型参数。
    """
    # 模拟变长数据
    # 生成一个随机整数，代表源序列的时间步长 T1，范围在 [300, 400) 之间
    # 生成一个形状为 (T1, 8) 的张量，填充随机正态分布的值。模拟长度T1行，8列的传感器数据，一个矩阵
    # 生成一个随机整数，代表时间步长 T2，范围在 [100, 150) 之间
    # 生成一个形状为 (T2, 1) 的张量，一列时间长度为T2的扭矩值
    # dataset = [
    #     (torch.randn(torch.randint(300, 400, (1,)).item(), 8), 
    #      torch.randn(torch.randint(100, 150, (1,)).item(), 1)) #一个元组，内容为 (传感器src, 扭矩标签tgt)，一条数据
    #     for _ in range(10000) #包含100个元组的列表，100条数据
    # ]
    # 导入真实的数据
    dataset = torch.load('dataset_tensor.pt',weights_only=False)
    print("数据类型:", type(dataset))
    print("样本数量:", len(dataset))
    # 取一个样本（可修改索引）
    sample = dataset[250]
    print("样本类型:", type(sample))
    print("样本维度:", len(sample))
    
    # 创建 DataLoader
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, collate_fn=collate_fn)
    print("dataloader: ", dataloader)  #<torch.utils.data.dataloader.DataLoader object at 0x7f8cd58bb8e0>
    # 初始化模型、损失函数和优化器
    model = GaitCNNTransformer(input_dim=9)
    # print("model: ", model) #打印模型的结构
    # 方法 1：简单保存
    with open("model_architecture.txt", "w") as f:
        print(model, file=f)

    criterion = nn.MSELoss()

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) #model.parameters()整个参数
    first_time=0
    # 训练并记录 loss
    loss_history = []  # 用于存储每个 epoch 的 loss
    # 训练循环
    for src, tgt, src_mask, tgt_mask in dataloader:
        optimizer.zero_grad()
        '''
        src :形状: (B, T1, 8)
        含义: 输入源序列， B 是批次大小， T1 是动态生成的序列长度， 8 是输入特征维度。
        tgt :形状: (B, T2, 1)
        含义: 目标序列（训练时使用）， T2 是动态生成的序列长度， 1 是输出特征维度。
        src_key_padding_mask :形状: (B, T1')
        含义: 源序列的填充掩码， T1' 是经过 CNN 降采样后的时间步长（ T1/4 ）。 True 表示填充部分。
        tgt_key_padding_mask :形状: (B, T2)
        含义: 目标序列的填充掩码， True 表示填充部分。
        forward(self, src, tgt=None, src_key_padding_mask=None, tgt_key_padding_mask=None)
        '''
        output = model.__call__(src, tgt, src_key_padding_mask=src_mask, tgt_key_padding_mask=tgt_mask)
        if first_time==0:
            print("type(output): ", type(output))
            # 如果是 Tensor，检查形状、数据类型、设备
            if isinstance(output, torch.Tensor):
                print("Shape:", output.shape)  # 例如 torch.Size([batch_size, seq_len, hidden_dim])
                print("Dtype:", output.dtype)  # 例如 torch.float32
                print("Device:", output.device)  # 例如 cuda:0 或 cpu
                print("Requires_grad:", output.requires_grad)  # 是否参与梯度计算
                # print("output: ", output[0:1])#取出一个数据,并打印内容
                print("output[3].shape", output[0:1].shape)   #torch.Size([1, 136, 1])  长度为136的一列(生成的扭矩曲线)
                #output[0:1].shape torch.Size([1, 149, 1]) 长度为149一行，output共有8行
                print("tgt.shape: ", tgt.shape) #torch.Size([8, 147, 1]) 8：batch, 长度为147一列（扭矩曲线）
            first_time=1
        
        # 只计算非 padding 部分的 loss
        loss = criterion(output[~tgt_mask.unsqueeze(-1)], tgt[~tgt_mask.unsqueeze(-1)])
        loss.backward()
        optimizer.step()
        loss_history.append(loss.item())
        if len(loss_history)%10==0:
            print(f"{len(loss_history)}, Loss: {loss.item():.4f}")


# 4. 保存 loss 为 CSV 文件
import   pandas as pd
loss_df = pd.DataFrame({"Epoch": range(1,len(loss_history)+1), "Loss": loss_history})
loss_df.to_csv("training_loss.csv", index=False)  # 不保存行索引

# 保存训练好的模型为独立文件
torch.save(model.state_dict(), 'gait_cnn_transformer.pth')
print("✅ 模型已保存为 gait_cnn_transformer.pth")